# -*- coding: utf-8 -*-
'''
Created on 2017年5月17日

@author: ZhuJiahui
'''

import os
import time
import re
from file_utils.file_writer import quick_write_1d_to_text

def endbpedia_get_triples(read_filename, write_directory):
    
    filtered_relations = ['wikiPage', 'language', 'contributorID', 'rdf-syntax', 'rdf-schema', 'contributor', 'isPrimaryTopicOf', 'primaryTopic', 'homepage', 'wgs84_pos']
    p1 = re.compile(r'(?<=/)([\s\S]*?)(?=\>)')
    p2 = re.compile(r'(?<=")([\s\S]*?)(?="\^\^)')
    
    segment_triple_text = []
    segment = 1000000
    file_count = 1
    
    with open(read_filename, 'r', encoding='utf-8') as f:
        for each_line in f:
            pre_matches = p1.findall(each_line.strip())
            if (len(pre_matches) == 3):
                this_relation = pre_matches[1].strip().split('/')[-1]
                filtered_flag = False
                for each in filtered_relations:
                    if each in this_relation:
                        filtered_flag = True
                        break
                if not filtered_flag:
                    this_head = pre_matches[0].strip().split('/')[-1]
                    this_tail = pre_matches[2].strip().split('/')[-1]
                    
                    if ('XMLSchema#' in this_tail):
                        this_matches = p2.findall(each_line.strip())
                        if (len(this_matches) > 0):
                            this_tail = this_matches[0].strip()
                    
                    this_triple_text = this_head + '\t|\t' + this_relation + '\t|\t' + this_tail
                    segment_triple_text.append(this_triple_text)
                    
                    if len(segment_triple_text) >= segment:
                        quick_write_1d_to_text(write_directory + '/' + str(file_count) + '.txt', segment_triple_text)
                        segment_triple_text = []
                        file_count += 1
    
    quick_write_1d_to_text(write_directory + '/' + str(file_count) + '.txt', segment_triple_text)

if __name__ == '__main__':
    start = time.clock()
    now_directory = os.getcwd()
    root_directory = os.path.dirname(now_directory) + '/'
    read_filename = 'D:/dbpedia_2016_09_26.nt'
    write_directory = root_directory + 'dataset/20newsgroup/dbpedia_triples'
    if (not(os.path.exists(write_directory))):
        os.mkdir(write_directory)
    
    endbpedia_get_triples(read_filename, write_directory)
    
    print('Total time %f seconds' % (time.clock() - start))
    